package com.lxg.ai.sse.impl;

import com.lxg.ai.sse.fallback.ChatCallback;
import org.jetbrains.annotations.NotNull;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

public class SseChatCallback implements ChatCallback {

    private StringBuffer stringBuffer = new StringBuffer();

    private final SseEmitter emitter;

    public SseChatCallback(SseEmitter emitter) {
        this.emitter = emitter;
    }

    @Override
    public void receivedMessagePart(@NotNull String messagePart) throws Exception {
        emitter.onError(throwable -> {
            throw new RuntimeException("SSE链接中断");
        });
        emitter.send(SseEmitter.event().data(messagePart));
        stringBuffer.append(messagePart);
    }

    @Override
    public void completed(@NotNull String message)  {
        System.out.println("完成了");
        emitter.complete();
        System.out.println(stringBuffer.toString());
    }

    @Override
    public void error(@NotNull String errMsg) {
        emitter.completeWithError(new Throwable(errMsg));
    }
}
